from fastapi import FastAPI, Request, Depends
from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine, AsyncSession, async_sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from typing import AsyncGenerator
from typing_extensions import Annotated
from config.app import get_config
from core.curd import Base

import logging
import os
import time

POOL_SIZE = 20
POOL_RECYCLE = 3600
POOL_TIMEOUT = 15
MAX_OVERFLOW = 2
CONNECT_TIMEOUT = 60
connect_args = {"connect_timeout": CONNECT_TIMEOUT}

DB_USERNAME = os.getenv('DB_USER', '127.0.0.1')
DB_PASSWORD = os.getenv('DB_PASSWORD', '123123')
DB_HOST = os.getenv('DB_HOST', '127.0.0.1')
DB_PORT = int(os.getenv('DB_PORT', 3307))
DB_DATABASE = os.getenv('DB_NAME', 'test')


class SQLFileHandler(logging.StreamHandler):
    def __init__(self, filename=''):
        super().__init__()
        self._filename = filename

    def emit(self, record):
        today = time.strftime("%Y%m%d")
        with open(get_config().ROOT_PATH + '/runtime/log/query_' + today + '.log', 'a') as f:
            msg = self.format(record) + '\n'
            f.write(msg)


async def register_mysql(app: FastAPI):
    try:
        async_engine: AsyncEngine = create_async_engine(
            'mysql+aiomysql://' + DB_USERNAME + ':' + DB_PASSWORD + '@' + DB_HOST + ':' + str(
                DB_PORT) + '/' + DB_DATABASE + '?charset=utf8mb4',
            pool_size=POOL_SIZE,
            pool_recycle=POOL_RECYCLE,
            pool_timeout=POOL_TIMEOUT,
            max_overflow=MAX_OVERFLOW,
            connect_args=connect_args,
            pool_pre_ping=True,
            pool_use_lifo=True,
            echo_pool=True,
            echo=False,
            isolation_level="READ UNCOMMITTED"  # 设置隔离级别参数isolation_level防止缓存后更新后结果查询不到的问题
        )

        logger = logging.getLogger(__name__)
        handler = SQLFileHandler()
        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
        handler.setLevel(logging.INFO)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        # 设置全局默认的日志记录器为上述自定义的logger
        logging.root.handlers = []
        logging.basicConfig(level=logging.DEBUG, handlers=[handler])

        app.state.db_async_session = async_sessionmaker(bind=async_engine, autoflush=True, expire_on_commit=True)
        app.state.db_async_engine = async_engine

        async with app.state.db_async_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

    except Exception as e:
        print(str(e))
        pass


async def dispose_mysql(app: FastAPI) -> None:
    await app.state.db_async_engine.dispose()


async def get_db_session(request: Request) -> AsyncGenerator[AsyncSession, None]:
    """
    Create and get database session.

    # :param request: current request.
    :yield: database session.
    """

    session: AsyncSession = AsyncSession(request.app.state.db_async_engine)

    try:  # noqa: WPS501
        yield session
    finally:
        await session.commit()
        await session.close()

    # async_session = async_sessionmaker(bind=app.state.db_async_engine, autoflush=True, expire_on_commit=True)
    #
    # try:
    #     yield async_session
    # except Exception as se:
    #     await async_session.rollback()
    #     raise se
    # finally:
    #     await async_session.close()

# Session Annotated
# CurrentSession = Annotated[AsyncSession, Depends(get_db_session())]


async def init_create_table(app):
    """
    应用启动时初始化数据库连接
    :return:
    """
    print("~~~~~~~~~~~~~~~~~~~~~~初始化数据库~~~~~~~~~~~~~~~~~~~~~~")
    async with app.state.db_async_engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)